# This code is hosted on http://code.google.com/p/lenthorp/
# Freely available for use in applications, but should NOT be modified
# Email all comments to lenthorpresearch@gmail.com

from quantlib import *


# Class that aids in generating multiple correlated paths for a variety of models
#   Note that any correlations within the macro processes should be zero, since these are set here by setCorrelation()
class StochasticPathGenerator():

    def __init__(self):
        # Dictionary of macro processes - each one can contain multiple processes - do not alter directly
        self._processes = {}

        self._miniProcesses = []

        # Dictionary of correlations - do not alter directly
        self._correlations = {}
        
        # List of available RNGs
        self.rngs = {'MersenneTwister-InverseNormal' : (lambda dim : GetGenericMersenneTwisterInverseNormal(dim,42))}
        self.chosenRng = ''

        # time grid of points to simulate at
        self._timeGrid = []

        self._initialised = False
        self._generator = 0
        self._ordering = []
        self._multiPath = 0

    # Only need to call this when something has changed
    def initialiseSimulation(self, timeGrid):
        self._timeGrid = timeGrid
        tempVec = DoubleVector(len(self._timeGrid))
        for idx in range(len(self._timeGrid)):
            tempVec[idx] = self._timeGrid[idx]
        qlTimeGrid = CreateTimeGridFromVector(tempVec)
        
        # make vector of processes
        nprocs = 0
        for proc in self._processes.values():
            nprocs += proc.size()
        nMacroProcs = len(self._processes.values())
        spspv = StochasticProcessSharedPointerVector(nMacroProcs)
        self._ordering = ['' for idx in range(nprocs)]
        count = 0
        for idx in range(nMacroProcs):
            spspv[idx] = self._processes.values()[idx]
            for jdx in range(self._processes.values()[idx].size()):
                self._ordering[count] = self._processes.keys()[idx] + '_' + str(jdx)
                count += 1

        # construct correlation matrix
        cm = Matrix(nprocs,nprocs)
        corrs = DoubleVector(nprocs * nprocs)
        for idx in range(nprocs):
            firstItem = self._ordering[idx]
            for jdx in range(idx, nprocs):
                secondItem = self._ordering[jdx]
                if idx == jdx:
                    val = 1.0
                else:
                    if firstItem + ',' + secondItem in self._correlations:
                        val = self._correlations[firstItem + ',' + secondItem]
                    else:
                        val = self._correlations[secondItem + ',' + firstItem]
                corrs[nprocs * idx + jdx] = val
                corrs[nprocs * jdx + idx] = val
        FillRowMajorMatrixFromVector(cm, corrs)
        
        stochasticProcessArray = GeneralStochasticProcessArray(spspv, cm)
        
        ndims = len(self._timeGrid) * nprocs
        rng = self.rngs[self.chosenRng](ndims)
        self._generator = MersenneInverseNormalMultiGenerator(stochasticProcessArray, qlTimeGrid, rng)
        self._initialised = True
        
    # Adds a macro process to the list to be generated, returns a 0 if successful
    def registerMacroProcess(self, macroProcess, macroProcessName):
        self._initialised = False
        self._processes[macroProcessName] = macroProcess
        nProcs = macroProcess.size()
        for procIdx in range(nProcs):
            for idx in range(len(self._miniProcesses)):
                self._correlations[macroProcessName + '_' + str(procIdx) + ',' + self._miniProcesses[idx]] = 0.0
            self._miniProcesses.append(macroProcessName + '_' + str(procIdx))

    def removeMacroProcess(self, macroProcessName):
        self._initialised = False

    def getPossibleGaussianRandomNumberGenerators(self):
        return self.rngs
    
    def setGaussianRandomNumberGenerator(self, rngName):
        if rngName in self.rngs:
            self.chosenRng = rngName
            self._initialised = False
        else:
            raise Exception('Given rngName not present')

    def setCorrelation(self, name1, index1, name2, index2, corr):
        self._initialised = False
        if name1 + '_' + str(index1) + ',' + name2 + '_' + str(index2) in self._correlations:
            self._correlations[name1 + '_' + str(index1) + ',' + name2 + '_' + str(index2)] = corr
        elif name2 + '_' + str(index2) + ',' + name1 + '_' + str(index1) in self._correlations:
            self._correlations[name2 + '_' + str(index2) + ',' + name1 + '_' + str(index1)] = corr
        else:
            raise Exception('Given names/indices do not exist')
        
    def simulateNext(self):
        if self._initialised:
            self._multiPath = self._generator.next()
        else:
            raise Exception('Generator needs to be initialised first')

    def getPath(self, name, index):
        if self._multiPath == 0:
            raise Exception('Generator needs to be simulated first')
        else:
            array = [0.0 for idx in range(len(self._timeGrid)+1)]
            path = self._multiPath.value.at(self._ordering.index(name + '_' + str(index)))
            for idx in range(len(self._timeGrid)+1):
                array[idx] = path.at(idx)
            return array